import torch
import networkx as nx

def calculate_adjacency_matrix(edge_list):
    node_set = set()
    for edge in edge_list:
        node_set.add(edge[0])
        node_set.add(edge[1])
    G = nx.DiGraph()
    G.add_nodes_from(node_set)
    for edge in edge_list:
        G.add_edge(edge[0], edge[1])
    return nx.adjacency_matrix(G).toarray()


def calculate_laplacian_with_self_loop(matrix):
    # TODO
    matrix = matrix + torch.eye(matrix.size(0))
    row_sum = matrix.sum(1)
    d_inv_sqrt = torch.pow(row_sum, -0.5).flatten()
    d_inv_sqrt[torch.isinf(d_inv_sqrt)] = 0.0
    d_mat_inv_sqrt = torch.diag(d_inv_sqrt)
    normalized_laplacian = (
        matrix.matmul(d_mat_inv_sqrt).transpose(0, 1).matmul(d_mat_inv_sqrt)
    )
    return normalized_laplacian
